package com.example.stream;

import cn.hutool.json.JSONUtil;
import com.example.entity.KafkaRecord;
import com.example.entity.TPUser;
import com.example.entity.User;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.function.FlatMapFunction;
import org.apache.spark.api.java.function.MapFunction;
import org.apache.spark.sql.*;
import org.apache.spark.sql.streaming.StreamingQuery;
import org.apache.spark.sql.streaming.StreamingQueryException;
import org.apache.spark.sql.streaming.Trigger;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.types.StructField;
import org.apache.spark.sql.types.StructType;

import java.sql.Timestamp;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.TimeUnit;

public class StructedStreamingDemo {
    private String ch_path = "spark-warehouse";

    public static void main(String[] args) {
        StructedStreamingDemo demo = new StructedStreamingDemo();
        demo.run13();
    }

    public void run13(){
        SparkSession sparkSession = SparkSession.builder().appName("structured-kafka").master("local").getOrCreate();
        sparkSession.sparkContext().setLogLevel("WARN");
        Dataset<Row> df = sparkSession.readStream()
                .format("kafka")
                .option("kafka.bootstrap.servers", "localhost:9092")
                .option("subscribe", "kafka-3")
                .load();
        Dataset<KafkaRecord> dataset = df.selectExpr("CAST(key AS STRING)", "CAST(value AS STRING)", "topic", "partition", "offset", "timestamp", "timestampType")
                .as(Encoders.bean(KafkaRecord.class));
        Dataset<User> userDataset = dataset.map((MapFunction<KafkaRecord, User>) x -> {
            String userjson = x.getValue();
            User user = JSONUtil.toBean(userjson, User.class);
            return user;
        }, Encoders.bean(User.class));

        StreamingQuery query = userDataset
                .writeStream()
                .format("console")
                .trigger(Trigger.ProcessingTime(5, TimeUnit.SECONDS))
                .start();
        try {
            query.awaitTermination();
        } catch (StreamingQueryException e) {
            e.printStackTrace();
        }
    }

    public void run12(){
        SparkSession sparkSession = SparkSession.builder().appName("structured-kafka").master("local").getOrCreate();
        sparkSession.sparkContext().setLogLevel("WARN");
        Dataset<Row> df = sparkSession.readStream()
                .format("kafka")
                .option("kafka.bootstrap.servers", "localhost:9092")
                .option("subscribe", "kafka-3")
                .load();
        Dataset<KafkaRecord> dataset = df.selectExpr("CAST(key AS STRING)", "CAST(value AS STRING)", "topic", "partition", "offset", "timestamp", "timestampType")
                .as(Encoders.bean(KafkaRecord.class));
        Dataset<User> userDataset = dataset.map((MapFunction<KafkaRecord, User>) x -> {
            String userjson = x.getValue();
            User user = JSONUtil.toBean(userjson, User.class);
            return user;
        }, Encoders.bean(User.class));

        StreamingQuery query = userDataset
               .writeStream()
                .foreach(new MysqlForeachWriter<User>())
                .trigger(Trigger.ProcessingTime(5, TimeUnit.SECONDS))
                .start();
        try {
            query.awaitTermination();
        } catch (StreamingQueryException e) {
            e.printStackTrace();
        }
    }

    public void run11(){
        SparkSession sparkSession = SparkSession.builder().appName("structured-kafka").master("local").getOrCreate();
        sparkSession.sparkContext().setLogLevel("WARN");
        Dataset<Row> df = sparkSession.readStream()
                .format("kafka")
                .option("kafka.bootstrap.servers", "localhost:9092")
                .option("subscribe", "kafka-3")
                .load();
        Dataset<KafkaRecord> dataset = df.selectExpr("CAST(key AS STRING)", "CAST(value AS STRING)", "topic", "partition", "offset", "timestamp", "timestampType")
                .as(Encoders.bean(KafkaRecord.class));
        Dataset<User> userDataset = dataset.map((MapFunction<KafkaRecord, User>) x -> {
            String userjson = x.getValue();
            User user = JSONUtil.toBean(userjson, User.class);
            return user;
        }, Encoders.bean(User.class));

//        Dataset<TUser> filter = userDataset.filter((FilterFunction<TUser>) value -> value.getAge() > 0);
//        Dataset<String> stringDataset = filter.map((MapFunction<TUser, String>) x -> x.getName(), Encoders.STRING());

//        Dataset<Row> rowDataset = userDataset.map((MapFunction<TUser, Tuple2<String, Integer>>) x -> {
//            Tuple2<String, Integer> tuple2 = new Tuple2<>(x.getName(), x.getAge());
//            return tuple2;
//        }, Encoders.tuple(Encoders.STRING(), Encoders.INT())).toDF("name", "age");

        StreamingQuery query = userDataset.writeStream()
                .format("console")
                .trigger(Trigger.ProcessingTime(5, TimeUnit.SECONDS))
                .start();

        try {
            query.awaitTermination();
        } catch (StreamingQueryException e) {
            e.printStackTrace();
        }
    }

    public void run10(){
        SparkSession sparkSession = SparkSession.builder().appName("structed-kafka").master("local").getOrCreate();
        sparkSession.sparkContext().setLogLevel("WARN");
        Dataset<Row> df = sparkSession
                .readStream()
                .format("kafka")
                .option("kafka.bootstrap.servers", "localhost:9092")
                .option("subscribe", "kafka-3")
                .load();
        Dataset<KafkaRecord> dataset = df.selectExpr("CAST(key AS STRING)", "CAST(value AS STRING)", "topic", "partition", "offset", "timestamp", "timestampType")
                .as(Encoders.bean(KafkaRecord.class));

        Dataset<User> userDataset = dataset.map((MapFunction<KafkaRecord, User>) x -> {
            String userjson = x.getValue();
            User user = JSONUtil.toBean(userjson, User.class);
            return user;
        }, Encoders.bean(User.class));

        Dataset<Row> rowDataset = userDataset
                .dropDuplicates("name", "timestamp")
                .withWatermark("timestamp", "5 minutes")
                .groupBy(
                        functions.window(userDataset.col("timestamp"), "10 minutes", "5 minutes"),
                        userDataset.col("name"),
                        userDataset.col("email")
                ).count();

        StreamingQuery query = rowDataset.writeStream()
                .outputMode("update") //  使用 complete 模式 必须结合聚合函数
                .format("console")
                .trigger(Trigger.ProcessingTime(5, TimeUnit.SECONDS))
                .start();

        System.out.printf("id:{%s}, name:{%s},runId:{%s}\n", query.id(), query.name(), query.runId());
        try {
            query.awaitTermination();
        } catch (StreamingQueryException e) {
            e.printStackTrace();
        }
    }

    public void run9(){
        SparkConf sparkConf = new SparkConf().setAppName("structed-kafka").setMaster("local");
        sparkConf.set("spark.serializer", "org.apache.spark.serializer.KryoSerializer");
        SparkSession sparkSession = SparkSession.builder().appName("structed-kafka").master("local").getOrCreate();
        sparkSession.sparkContext().setLogLevel("WARN");
        Dataset<Row> df = sparkSession
                .readStream()
                .format("kafka")
                .option("kafka.bootstrap.servers", "localhost:9092")
                .option("subscribe", "kafka-3")
                .load();
        Dataset<KafkaRecord> dataset = df.selectExpr("CAST(key AS STRING)", "CAST(value AS STRING)", "topic", "partition", "offset", "timestamp", "timestampType")
                .as(Encoders.bean(KafkaRecord.class));

        Dataset<User> userDataset = dataset.map((MapFunction<KafkaRecord, User>) x -> {
            Timestamp timestamp = new Timestamp(x.getTimestamp());
            String userjson = x.getValue();
            User user = JSONUtil.toBean(userjson, User.class);
            User tUser = new User(user.getId(), user.getName(), user.getAge(), user.getEmail(), timestamp);
            return tUser;
        }, Encoders.bean(User.class));

        StreamingQuery query = userDataset.writeStream()
                .outputMode("append") //  使用 complete 模式 必须结合聚合函数
                .trigger(Trigger.ProcessingTime(5, TimeUnit.SECONDS))
                .format("console")
                .start();
        try {
            query.awaitTermination();
        } catch (StreamingQueryException e) {
            e.printStackTrace();
        }
    }

    public void run8(){
        SparkConf sparkConf = new SparkConf().setAppName("structed-kafka").setMaster("local");
        sparkConf.set("spark.serializer", "org.apache.spark.serializer.KryoSerializer");
        SparkSession sparkSession = SparkSession.builder().appName("structed-kafka").master("local").getOrCreate();
        sparkSession.sparkContext().setLogLevel("WARN");
        Dataset<Row> df = sparkSession
                .readStream()
                .format("kafka")
                .option("kafka.bootstrap.servers", "localhost:9092")
                .option("subscribe", "kafka-3")
                .load();

        Dataset<KafkaRecord> dataset = df.selectExpr("CAST(key AS STRING)", "CAST(value AS STRING)", "topic", "partition", "offset", "timestamp", "timestampType")
                .as(Encoders.bean(KafkaRecord.class));

        Dataset<Row> rowDataset = dataset
                .withWatermark("timestamp", "5 minutes")
                .groupBy(
                        functions.window(dataset.col("timestamp"), "2 minutes", "2 minutes"),
                        dataset.col("partition")
                ).count();

        StreamingQuery query = rowDataset.writeStream()
                .outputMode("append") //  complete
                .trigger(Trigger.ProcessingTime(1, TimeUnit.SECONDS))
                .format("console")
                .start();
        try {
            query.awaitTermination();
        } catch (StreamingQueryException e) {
            e.printStackTrace();
        }
    }

    public void run7(){
        SparkConf sparkConf = new SparkConf().setAppName("structed-kafka").setMaster("local");
        sparkConf.set("spark.serializer", "org.apache.spark.serializer.KryoSerializer");
        SparkSession sparkSession = SparkSession.builder().appName("structed-kafka").master("local").getOrCreate();
        sparkSession.sparkContext().setLogLevel("WARN");
        Dataset<Row> df = sparkSession
                .readStream()
                .format("kafka")
                .option("kafka.bootstrap.servers", "localhost:9092")
                .option("subscribe", "kafka-3")
                .load();

        Dataset<Row> rowDataset = df.groupBy(df.col("partition")).count();

        StreamingQuery query = rowDataset.writeStream()
                .outputMode("complete") // append
                .trigger(Trigger.ProcessingTime(5, TimeUnit.SECONDS))
                .format("console")
                .start();
        try {
            query.awaitTermination();
        } catch (StreamingQueryException e) {
            e.printStackTrace();
        }
    }



    public void run5(){
        SparkConf sparkConf = new SparkConf().setAppName("structed-kafka").setMaster("local");
        sparkConf.set("spark.serializer", "org.apache.spark.serializer.KryoSerializer");
        SparkSession sparkSession = SparkSession.builder().appName("structed-kafka").master("local").getOrCreate();
        sparkSession.sparkContext().setLogLevel("WARN");
        Dataset<Row> df = sparkSession
                .readStream()
                .format("kafka")
                .option("kafka.bootstrap.servers", "localhost:9092")
                .option("subscribe", "kafka-3")
                .load();
        Dataset<TPUser> dataset = df.selectExpr("CAST(key AS STRING)", "CAST(value AS STRING)").as(Encoders.bean(TPUser.class));
        Dataset<User> userDataset = dataset.map((MapFunction<TPUser, User>) x -> {
            User user = JSONUtil.toBean(x.getValue(), User.class);
            return user;
        }, Encoders.bean(User.class));

        StreamingQuery query = userDataset.writeStream().format("console").start();
        try {
            query.awaitTermination();
        } catch (StreamingQueryException e) {
            e.printStackTrace();
        }
    }

    public void run4(){
        SparkConf sparkConf = new SparkConf().setAppName("structed-kafka").setMaster("local");
        sparkConf.set("spark.serializer", "org.apache.spark.serializer.KryoSerializer");
        SparkSession sparkSession = SparkSession.builder().appName("structed-kafka").master("local").getOrCreate();
        sparkSession.sparkContext().setLogLevel("WARN");
        Dataset<Row> df = sparkSession
                .readStream()
                .format("kafka")
                .option("kafka.bootstrap.servers", "localhost:9092")
                .option("subscribe", "kafka-3")
                .load();
        Dataset<String> dataset = df.selectExpr("CAST(value AS STRING)").as(Encoders.bean(String.class));
//        Dataset<User> userDataset = dataset.flatMap((FlatMapFunction<String, User>) x -> {
//            User user = JSONUtil.toBean(x, User.class);
//            return Arrays.asList(user).iterator();
//        }, Encoders.bean(User.class));
        Dataset<User> userDataset = dataset.map((MapFunction<String, User>) x -> {
            User user = JSONUtil.toBean(x, User.class);
            return user;
        }, Encoders.bean(User.class));

        StreamingQuery query = dataset.writeStream().format("console").start();
        try {
            query.awaitTermination();
        } catch (StreamingQueryException e) {
            e.printStackTrace();
        }

    }

    public void run3(){
        SparkConf sparkConf = new SparkConf().setAppName("structed-kafka").setMaster("local");
        sparkConf.set("spark.serializer", "org.apache.spark.serializer.KryoSerializer");
        SparkSession sparkSession = SparkSession.builder().appName("structed-kafka").master("local").getOrCreate();
        sparkSession.sparkContext().setLogLevel("WARN");
        Dataset<Row> df = sparkSession
                .readStream()
                .format("kafka")
                .option("kafka.bootstrap.servers", "localhost:9092")
                .option("subscribe", "kafka-3")
                .load();
        Dataset<String> dataset = df.selectExpr("CAST(value AS STRING)").as(Encoders.STRING());

        StreamingQuery query = dataset.writeStream()
                .option("checkpointLocation", ch_path)
//                .format("console")
                .format("kafka")
                .option("kafka.bootstrap.servers", "localhost:9092")
                .option("topic", "kafka-2")
                .start();
        try {
            query.awaitTermination();
        } catch (StreamingQueryException e) {
            e.printStackTrace();
        }
    }

    public void run(){
        SparkConf sparkConf = new SparkConf().setAppName("structed-kafka").setMaster("local");
        sparkConf.set("spark.serializer", "org.apache.spark.serializer.KryoSerializer");
        SparkSession sparkSession = SparkSession.builder().appName("structed-kafka").master("local").getOrCreate();
        Dataset<Row> df = sparkSession
                .readStream()
                .format("kafka")
                .option("kafka.bootstrap.servers", "localhost:9092")
                .option("subscribe", "kafka-3")
                .load();
//        Dataset<TPUser> dataset = df.selectExpr("CAST(timestamp AS STRING)", "CAST(value AS STRING)").as(Encoders.bean(TPUser.class));
        Dataset<String> dataset = df.selectExpr("CAST(value AS STRING)").as(Encoders.STRING());
        Dataset<String> stringDataset = dataset.flatMap((FlatMapFunction<String, String>) x -> Arrays.asList(x).iterator(), Encoders.STRING());

        StreamingQuery query = stringDataset.writeStream().trigger(Trigger.ProcessingTime(2)).format("console").start();
        try {
            query.awaitTermination();
        } catch (StreamingQueryException e) {
            e.printStackTrace();
        }
/*
+----+--------------------+-------+---------+------+--------------------+-------------+
| key|               value|  topic|partition|offset|           timestamp|timestampType|
+----+--------------------+-------+---------+------+--------------------+-------------+
|null|[7B 22 6E 61 6D 6...|kafka-3|        0|    54|2021-01-29 09:12:...|            0|
|null|[7B 22 6E 61 6D 6...|kafka-3|        0|    55|2021-01-29 09:12:...|            0|
|null|[7B 22 6E 61 6D 6...|kafka-3|        2|    47|2021-01-29 09:12:...|            0|
|null|[7B 22 6E 61 6D 6...|kafka-3|        2|    48|2021-01-29 09:12:...|            0|
|null|[7B 22 6E 61 6D 6...|kafka-3|        1|    54|2021-01-29 09:12:...|            0|
+----+--------------------+-------+---------+------+--------------------+-------------+
*/

    }

    public void readFileStream(){
        SparkConf sparkConf = new SparkConf().setAppName("structed-kafka").setMaster("local");
        sparkConf.set("spark.serializer", "org.apache.spark.serializer.KryoSerializer");
        SparkSession sparkSession = SparkSession.builder().appName("structed-kafka").master("local").getOrCreate();
        sparkSession.sparkContext().setLogLevel("WARN");
        Dataset<Row> df = sparkSession
                .readStream()
                .schema(getSchema2())      // Specify schema of the csv files
                .csv("/var/stream-files");    // Equivalent to format("csv").load("/path/to/directory")

        StreamingQuery query = df.writeStream()
                .format("console")
                .start();
        try {
            query.awaitTermination();
        } catch (StreamingQueryException e) {
            e.printStackTrace();
        }
    }

    public void run2(){
        SparkSession sparkSession = SparkSession.builder().appName("structed-file").master("local").getOrCreate();
        sparkSession.sparkContext().setLogLevel("WARN");
        Dataset<Row> df = sparkSession
                .readStream()
                .schema(getSchema())
                .csv("/var/stream-files/");

        df.printSchema();
        StreamingQuery query = df.writeStream().format("console").start();
        try {
            query.awaitTermination();
        } catch (StreamingQueryException e) {
            e.printStackTrace();
        }

    }

    public StructType getSchema2(){
        List<StructField> fieldList = new ArrayList<>();
        StructField id = DataTypes.createStructField("id", DataTypes.IntegerType, true);
        StructField name = DataTypes.createStructField("name", DataTypes.StringType, true);

        fieldList.add(id);
        fieldList.add(name);
        StructType structType = DataTypes.createStructType(fieldList);
        return structType;
    }

    public StructType getSchema(){
        List<StructField> fieldList = new ArrayList<>();
        StructField id = DataTypes.createStructField("id", DataTypes.IntegerType, true);
        StructField name = DataTypes.createStructField("name", DataTypes.StringType, true);
        StructField age = DataTypes.createStructField("age", DataTypes.IntegerType, true);
        StructField email = DataTypes.createStructField("email", DataTypes.IntegerType, true);
        StructField timestamp = DataTypes.createStructField("timestamp", DataTypes.TimestampType, true);
        fieldList.add(id);
        fieldList.add(name);
        fieldList.add(age);
        fieldList.add(email);
        fieldList.add(timestamp);
        StructType structType = DataTypes.createStructType(fieldList);
        return structType;
    }

}